# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import datetime
import logging
import tabulate

from coin.base.param_util import to_list


class SpotBalancePrinter(object):
  def __init__(self, products, bookmap, reservemap, currency=None, logger=None):
    self._logger = logger or logging.getLogger('SpotBalancePrinter')
    self._products = sorted(to_list(products))
    self._bookmap = bookmap
    self._reservemap = reservemap
    currency = currency or 'USDT'
    if isinstance(currency, str):
      self._currency_str = currency
    else:
      self._currency_str = str(currency.currency)

    currencies = []
    for product in self._products:
      assert product.base.currency in reservemap, product
      assert product.quote.currency == self._currency_str
      currencies.append(product.base.currency)
    assert len(set(currencies)) == len(currencies), currencies
    self._currencies = list(set(currencies))

    self._og = None
    self.started = False

  def register_order_gateway(self, order_gateway):
    self._og = order_gateway

  def start_print(self, ioloop, period_sec=10):
    if self.started:
      return
    self.started = True
    ioloop.add_timeout(datetime.timedelta(seconds=period_sec),
                       self._periodic_print,
                       ioloop,
                       period_sec)

  def _periodic_print(self, ioloop, period_sec):
    try:
      self.print_reserve_balance()
    finally:
      ioloop.add_timeout(datetime.timedelta(seconds=period_sec),
                         self._periodic_print,
                         ioloop,
                         period_sec)

  def _get_reserve_balance_price(self, product):
    reserve = self._reservemap.get(product.base.currency, None)
    balance = self._og.get_balance(product.base.currency)

    book = self._bookmap.get_book(product)
    if book is None or not book.has_ask() or not book.has_bid():
      mtm_price = None
    else:
      mtm_price = 0.5 * (book.ask0().price + book.bid0().price)
    return reserve, balance, mtm_price

  def print_reserve_balance(self, ignore_error=False):
    try:
      if self._og is None or not self._og.is_ready():
        return

      table = []
      total = 0.
      pnlbal_total = 0.
      for product in self._products:
        reserve, balance, price = self._get_reserve_balance_price(product)
        total += (balance or 0.)
        table.append([
            '*',
            product.base.currency,
            reserve,
            balance,
            price,
            balance * price, (balance - reserve) * price
        ])
        total += balance * price
        pnlbal_total += (balance - reserve) * price

      balance = self._og.get_balance(self._currency_str)
      total += balance
      pnlbal_total += balance

      table.append(['*', self._currency_str, None, balance, 1, balance, balance])
      table.append(['*', 'Total', None, None, None, total, pnlbal_total])

      pos_bal_table_str = tabulate.tabulate(table,
                                            headers=[
                                                '*',
                                                'Currency',
                                                'Reserve',
                                                'Balance',
                                                'QuotePrice',
                                                'BalanceInQuote',
                                                'PnlBalanceInQuote'
                                            ],
                                            floatfmt='.4f')
      self._logger.info('----- Reserve and balance:\n* %s\n%s',
                        str(datetime.datetime.now()),
                        pos_bal_table_str)
    except Exception:
      if ignore_error:
        self._logger.exception('Exception raised')
      else:
        raise
